import argparse
import cv2
import numpy as np
import torch

import sys
import os.path as osp
sys.path.append(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))))
from lzc.new_files.preprocess import preprocess
from models.experimental import attempt_load
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, scale_coords
from utils.torch_utils import select_device


def infer_pt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='xxx.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='xxx.jpg', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.35, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.65, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', default=True, help='display results')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    opt = parser.parse_args()

    opt.weights = '/weights/chuwei_best_20240805.pt'
    opt.source = '/data/chuwei/toilet_30.jpg'

    device = select_device(opt.device)
    model = attempt_load(opt.weights, map_location=device)  # load FP32 model
    model.eval()

    img = cv2.imread(opt.source)
    im, ratio, dwdh = preprocess(img)

    with torch.no_grad():  # Calculating gradients would cause a GPU memory leak
        input = torch.from_numpy(im).to(device)
        pred = model(input, augment=opt.augment)
        if isinstance(pred, tuple):
            pred = pred[0]
        # pred_np_bofore_nms = pred.cpu().numpy()
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        pred_np = pred[0].cpu().numpy()

    return pred_np


if __name__ == '__main__':
    pred = infer_pt()
    print('finish')
